package com.atguigu.hadoop.examples.mapreduce.terasort;

import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.List;
import java.util.zip.Checksum;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.conf.Configured;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.io.NullWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.io.WritableUtils;
import org.apache.hadoop.mapreduce.Counter;
import org.apache.hadoop.mapreduce.InputFormat;
import org.apache.hadoop.mapreduce.InputSplit;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.JobContext;
import org.apache.hadoop.mapreduce.Mapper;
import org.apache.hadoop.mapreduce.Mapper.Context;
import org.apache.hadoop.mapreduce.RecordReader;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
import org.apache.hadoop.util.PureJavaCrc32;
import org.apache.hadoop.util.Tool;
import org.apache.hadoop.util.ToolRunner;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class TeraGen extends Configured
        implements Tool
{
    private static final Logger LOG = LoggerFactory.getLogger(TeraGen.class);

    static long getNumberOfRows(JobContext job)
    {
        return job.getConfiguration().getLong(TeraSortConfigKeys.NUM_ROWS.key(), 0L);
    }

    static void setNumberOfRows(Job job, long numRows)
    {
        job.getConfiguration().setLong(TeraSortConfigKeys.NUM_ROWS.key(), numRows);
    }

    private static void usage()
            throws IOException
    {
        System.err.println("teragen <num rows> <output dir>");
        System.err.println("If you want to generate data and store them as erasure code striping file, just make sure that the parent dir of <output dir> has erasure code policy set");
    }

    private static long parseHumanLong(String str)
    {
        char tail = str.charAt(str.length() - 1);
        long base = 1L;
        switch (tail) {
            case 't':
                base *= -727379968L;
                break;
            case 'b':
                base *= 1000000000L;
                break;
            case 'm':
                base *= 1000000L;
                break;
            case 'k':
                base *= 1000L;
                break;
        }

        if (base != 1L) {
            str = str.substring(0, str.length() - 1);
        }
        return Long.parseLong(str) * base;
    }

    public int run(String[] args)
            throws IOException, InterruptedException, ClassNotFoundException
    {
        Job job = Job.getInstance(getConf());
        if (args.length != 2) {
            usage();
            return 2;
        }
        setNumberOfRows(job, parseHumanLong(args[0]));
        Path outputDir = new Path(args[1]);
        FileOutputFormat.setOutputPath(job, outputDir);
        job.setJobName("TeraGen");
        job.setJarByClass(TeraGen.class);
        job.setMapperClass(SortGenMapper.class);
        job.setNumReduceTasks(0);
        job.setOutputKeyClass(Text.class);
        job.setOutputValueClass(Text.class);
        job.setInputFormatClass(RangeInputFormat.class);
        job.setOutputFormatClass(TeraOutputFormat.class);
        return job.waitForCompletion(true) ? 0 : 1;
    }

    public static void main(String[] args) throws Exception {
        int res = ToolRunner.run(new Configuration(), new TeraGen(), args);
        System.exit(res);
    }

    public static class SortGenMapper extends Mapper<LongWritable, NullWritable, Text, Text>
    {
        private Text key = new Text();
        private Text value = new Text();
        private Unsigned16 rand = null;
        private Unsigned16 rowId = null;
        private Unsigned16 checksum = new Unsigned16();
        private Checksum crc32 = new PureJavaCrc32();
        private Unsigned16 total = new Unsigned16();
        private static final Unsigned16 ONE = new Unsigned16(1L);
        private byte[] buffer = new byte[100];
        private Counter checksumCounter;

        public void map(LongWritable row, NullWritable ignored, Mapper<LongWritable, NullWritable, Text, Text>.Context context)
                throws IOException, InterruptedException
        {
            if (this.rand == null) {
                this.rowId = new Unsigned16(row.get());
                this.rand = Random16.skipAhead(this.rowId);
                this.checksumCounter = context.getCounter(TeraGen.Counters.CHECKSUM);
            }
            Random16.nextRand(this.rand);
           // GenSort.generateRecord(this.buffer, this.rand, this.rowId);
            this.key.set(this.buffer, 0, 10);
            this.value.set(this.buffer, 10, 90);

            context.write(this.key, this.value);
            this.crc32.reset();
            this.crc32.update(this.buffer, 0, 100);

            this.checksum.set(this.crc32.getValue());
            this.total.add(this.checksum);
            this.rowId.add(ONE);
        }

        public void cleanup(Mapper<LongWritable, NullWritable, Text, Text>.Context context)
        {
            if (this.checksumCounter != null)
                this.checksumCounter.increment(this.total.getLow8());
        }
    }

    static class RangeInputFormat extends InputFormat<LongWritable, NullWritable>
    {
        public RecordReader<LongWritable, NullWritable> createRecordReader(InputSplit split, TaskAttemptContext context)
                throws IOException
        {
            return new RangeRecordReader();
        }

        public List<InputSplit> getSplits(JobContext job)
        {
            long totalRows = TeraGen.getNumberOfRows(job);
            int numSplits = job.getConfiguration().getInt("mapreduce.job.maps", 1);
            TeraGen.LOG.info("Generating " + totalRows + " using " + numSplits);
            List splits = new ArrayList();
            long currentRow = 0L;
            for (int split = 0; split < numSplits; split++)
            {
                long goal = (long)Math.ceil(totalRows * (split + 1) / numSplits);

                splits.add(new RangeInputSplit(currentRow, goal - currentRow));
                currentRow = goal;
            }
            return splits;
        }

        static class RangeRecordReader extends RecordReader<LongWritable, NullWritable>
        {
            long startRow;
            long finishedRows;
            long totalRows;
            LongWritable key = null;

            public void initialize(InputSplit split, TaskAttemptContext context)
                    throws IOException, InterruptedException
            {
                this.startRow = ((TeraGen.RangeInputFormat.RangeInputSplit)split).firstRow;
                this.finishedRows = 0L;
                this.totalRows = ((TeraGen.RangeInputFormat.RangeInputSplit)split).rowCount;
            }

            public void close() throws IOException
            {
            }

            public LongWritable getCurrentKey() {
                return this.key;
            }

            public NullWritable getCurrentValue() {
                return NullWritable.get();
            }

            public float getProgress() throws IOException {
                return (float)this.finishedRows / (float)this.totalRows;
            }

            public boolean nextKeyValue() {
                if (this.key == null) {
                    this.key = new LongWritable();
                }
                if (this.finishedRows < this.totalRows) {
                    this.key.set(this.startRow + this.finishedRows);
                    this.finishedRows += 1L;
                    return true;
                }
                return false;
            }
        }

        static class RangeInputSplit extends InputSplit
                implements Writable
        {
            long firstRow;
            long rowCount;

            public RangeInputSplit()
            {
            }

            public RangeInputSplit(long offset, long length)
            {
                this.firstRow = offset;
                this.rowCount = length;
            }

            public long getLength() throws IOException {
                return 0L;
            }

            public String[] getLocations() throws IOException {
                return new String[0];
            }

            public void readFields(DataInput in) throws IOException {
                this.firstRow = WritableUtils.readVLong(in);
                this.rowCount = WritableUtils.readVLong(in);
            }

            public void write(DataOutput out) throws IOException {
                WritableUtils.writeVLong(out, this.firstRow);
                WritableUtils.writeVLong(out, this.rowCount);
            }
        }
    }

    public static enum Counters
    {
        CHECKSUM;
    }
}